from typing import Dict, Set
from fastapi import WebSocket


class WebSocketHub:
    def __init__(self):
        self._subs: Dict[str, Set[WebSocket]] = {}

    async def subscribe(self, flight_no: str, ws: WebSocket):
        await ws.accept()
        s = self._subs.get(flight_no) or set()
        s.add(ws)
        self._subs[flight_no] = s

    def unsubscribe(self, flight_no: str, ws: WebSocket):
        s = self._subs.get(flight_no) or set()
        if ws in s:
            s.remove(ws)
        if not s:
            self._subs.pop(flight_no, None)

    async def broadcast(self, flight_no: str, payload: dict):
        s = list(self._subs.get(flight_no) or [])
        dead = []
        for ws in s:
            try:
                await ws.send_json(payload)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.unsubscribe(flight_no, ws)


hub = WebSocketHub()